"""
混合检索器
结合向量检索和关键词检索
"""

import logging
from typing import List, Dict, Any, Optional
from llama_index.core import VectorStoreIndex
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.core.schema import NodeWithScore

from ..models.llm_config import LLMConfig
from ..utils.logger import setup_logger
from ..utils.text_utils import extract_keywords

logger = setup_logger(__name__)


class HybridRetriever:
    """混合检索器"""
    
    def __init__(self, index: VectorStoreIndex, config: LLMConfig):
        """
        初始化混合检索器
        
        Args:
            index: 向量索引
            config: 配置对象
        """
        self.index = index
        self.config = config
        
        # 创建向量检索器
        self.vector_retriever = VectorIndexRetriever(
            index=index,
            similarity_top_k=config.similarity_top_k
        )
        
        # 创建相似度后处理器
        self.similarity_postprocessor = SimilarityPostprocessor(
            similarity_cutoff=0.3
        )
        
        logger.info("混合检索器初始化完成")
    
    def vector_search(self, query: str) -> List[NodeWithScore]:
        """
        向量检索
        
        Args:
            query: 查询文本
            
        Returns:
            检索结果节点列表
        """
        try:
            # 直接使用向量检索器
            nodes = self.vector_retriever.retrieve(query)
            
            # 应用相似度过滤
            filtered_nodes = self.similarity_postprocessor.postprocess_nodes(
                nodes, query
            )
            
            logger.info(f"向量检索返回 {len(filtered_nodes)} 个结果")
            return filtered_nodes
            
        except Exception as e:
            logger.error(f"向量检索失败: {e}")
            return []
    
    def keyword_search(self, query: str) -> List[NodeWithScore]:
        """
        关键词检索
        
        Args:
            query: 查询文本
            
        Returns:
            检索结果节点列表
        """
        try:
            # 提取关键词
            keywords = extract_keywords(query, min_length=2)
            if not keywords:
                logger.warning("未提取到有效关键词")
                return []
            
            # 构建关键词查询
            keyword_query = " ".join(keywords)
            
            # 使用向量检索器进行关键词检索
            nodes = self.vector_retriever.retrieve(keyword_query)
            
            logger.info(f"关键词检索返回 {len(nodes)} 个结果")
            return nodes
            
        except Exception as e:
            logger.error(f"关键词检索失败: {e}")
            return []
    
    def hybrid_search(self, query: str) -> List[NodeWithScore]:
        """
        混合检索
        
        Args:
            query: 查询文本
            
        Returns:
            混合检索结果
        """
        try:
            logger.info(f"开始混合检索: {query}")
            
            # 向量检索
            vector_results = self.vector_search(query)
            
            # 关键词检索
            keyword_results = self.keyword_search(query)
            
            # 合并结果
            all_results = vector_results + keyword_results
            
            # 去重（基于节点 ID）
            unique_results = {}
            for node in all_results:
                node_id = node.node.node_id
                if node_id not in unique_results:
                    unique_results[node_id] = node
                else:
                    # 保留分数更高的结果
                    if node.score > unique_results[node_id].score:
                        unique_results[node_id] = node
            
            # 转换为列表并排序
            final_results = list(unique_results.values())
            final_results.sort(key=lambda x: x.score, reverse=True)
            
            # 限制结果数量
            top_k = min(self.config.similarity_top_k, len(final_results))
            final_results = final_results[:top_k]
            
            logger.info(f"混合检索完成，返回 {len(final_results)} 个结果")
            return final_results
            
        except Exception as e:
            logger.error(f"混合检索失败: {e}")
            return []
    
    def search(self, query: str, search_type: str = "hybrid") -> List[NodeWithScore]:
        """
        执行检索
        
        Args:
            query: 查询文本
            search_type: 检索类型 ("vector", "keyword", "hybrid")
            
        Returns:
            检索结果
        """
        if search_type == "vector":
            return self.vector_search(query)
        elif search_type == "keyword":
            return self.keyword_search(query)
        elif search_type == "hybrid":
            return self.hybrid_search(query)
        else:
            logger.warning(f"未知的检索类型: {search_type}，使用混合检索")
            return self.hybrid_search(query)
